package edu.arizona.sista.learn.evaluation;

import java.io.BufferedReader;
import java.io.DataInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.InputStreamReader;
import java.text.DecimalFormat;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import edu.arizona.sista.learn.activity.model.ActivityInstance;
import edu.arizona.sista.learn.activity.model.PredicateInterval;

public class PerFramePRMetric {
	
	private static Set<String> _verbs = new HashSet<String>(
			Arrays.asList("approach", "chase", "flee", "follow",
					"movepos", "walk", "run", "go", "stop", "pass"));

	private static Set<String> _type = new HashSet<String>(
			Arrays.asList("person", "car", "truck", "bicycle", "motorbike"));
	
	class PRScore {
		public int tp, fp, tn, fn;
		PRScore() {
			this(0, 0, 0, 0);
		}
		PRScore(int tp, int fp, int tn, int fn) {
			this.tp = tp; this.fp = fp; this.tn = tn; this.fn = fn;
		}
		public double precision() {
		    if (tp == 0 && fp == 0)
		        return 1.0;
		    return (double)tp / (double)(tp + fp);
		}
		public double recall() {
		    if (tp == 0 && fn == 0)
		        return 1.0;
		    return (double)tp / (double)(tp + fn);
		}
		public double f1() {
		    double p = precision();
			double r = recall();
		    return 2. * ( p * r ) / ( p + r );
		}
		public double accuracy() {
		    if (tp == 0 && fp == 0 && tn == 0 && fn == 0) 
		        return -1.;
		    return (double)(tp + tn) / (double)(tp + fp + tn + fn);
		}
		public double relativeRisk() {
		    int a = tp, b = fp, c = fn, d = tn;
		    if (c == 0 || (a == 0 && b == 0))
		        return -1.;
		    return (double)(a*(c+d)) / (double)(c*(a+b));
		}
		public double falseRate() {
		    int a = tp, b = fp, c = fn, d = tn;
		    if (b == 0 || (c == 0 && d == 0))
		        return -1.;
		    return (double)(d*(a+b)) / (double)(b*(c+d));
		}
		public String toString() {
			DecimalFormat df = new DecimalFormat("#.00000");
			String str = tp + "\t" + fp + "\t" + fn + "\t" + tn +
					"\t" + df.format(precision()) + "\t" + df.format(recall()) + "\t" + df.format(f1()) +
					"\t" + df.format(relativeRisk()) + "\t" + df.format(falseRate());
			return str;
		}
	}
	
	private int toInteger(String str) {
		try {
			return Integer.parseInt(str);
		} catch (Exception e) {
			e.printStackTrace();
		}
		return Integer.MIN_VALUE;
	}

	public void evaluate(File goldDir, File predictionDir, int winSize, int stepSize, boolean typeRestriction) {
		if (!goldDir.isDirectory() || !predictionDir.isDirectory())
			return;
		
		// Load the gold
		Map<String, CMUXML> gold = loadGoldXMLs(goldDir, winSize, stepSize);
		
		// Load the predictions into timeline format
		Map<String, ActivityInstance> preds = loadPredictions(predictionDir);
		Map<String, List<Set<String>>> predTimeLines = new HashMap<String, List<Set<String>>>();
		for (String movie : preds.keySet()) {
			List<Set<String>> timeline = new ArrayList<Set<String>>();
			ActivityInstance inst = preds.get(movie);
			for (int i = 0; i <= inst.end(); i++) {
				timeline.add(null);
			}
			for (String key : inst.getIntervalsMap().keySet()) {
				List<PredicateInterval> intv = inst.getIntervals(key);
				for (PredicateInterval i : intv) {
					for (int t = i.start; t <= i.end; t++) { // NOTE: assuming the end is inclusive
						if (timeline.get(t) == null)
							timeline.set(t, new HashSet<String>());
						timeline.get(t).add(i.getPredicate().getName());
					}
				}
			}
			predTimeLines.put(movie, timeline);
		}
		
		// Initialize scores
		Map<String, PRScore> scores = new HashMap<String, PerFramePRMetric.PRScore>();
		for (String v : _verbs) {
			scores.put(v, new PRScore());
		}

		// Compare gold & preds timelines
		for (String movie : gold.keySet()) {
			CMUXML standard = gold.get(movie);

			String key = (movie.endsWith(".xml")) ? movie.substring(0, movie.length()-4) : movie;
			List<Set<String>> results = (predTimeLines.containsKey(key)) ?
					predTimeLines.get(key) : null;
			
			for (int i = 0; i < standard.getMovieLength(); i++) {
				Set<String> predictions = (results != null && results.size() > i) ?
						results.get(i) : null;
				for (String v : scores.keySet()) {
					boolean inGold = (typeRestriction) ? standard.containsAt(v, i, _type) : standard.containsAt(v, i);
					boolean inPred = (predictions != null) ? predictions.contains(v) : false;
					if (inGold && inPred)
						scores.get(v).tp++;
					else if (inPred)
						scores.get(v).fp++;
					else if (inGold)
						scores.get(v).fn++;
					else
						scores.get(v).tn++;
				}
			}
		}
		
		// Print out scores
		for (String v : scores.keySet()) {
			System.out.println(v + "\t" + scores.get(v).toString());
		}
	}
	
	private Map<String, ActivityInstance> loadPredictions(File predDir) {
		Map<String, ActivityInstance> preds = new HashMap<String, ActivityInstance>();
		if (predDir.isDirectory()) {
			for (File child : predDir.listFiles()) {
				if (child.getName().startsWith(".")) {
					continue;  // Ignore hidden files
				}

				File out = new File(child.getAbsolutePath() + "/activities.txt");
				ActivityInstance inst = new ActivityInstance();
				inst.setFilename(child.getName());

				try {
					DataInputStream in = new DataInputStream(new FileInputStream(out));
					BufferedReader br = new BufferedReader(new InputStreamReader(in));
					String strLine;
					while ((strLine = br.readLine()) != null)   {
						String[] parts = strLine.trim().split("\\s+");
						String label = parts[0].split("\\(")[0];
						int start = toInteger(parts[1]);
						int end = toInteger(parts[2]);
						PredicateInterval i = PredicateInterval.make(label, start, end);
						inst.addInterval(i);
					}
					in.close();
				}catch (Exception e) {
					e.printStackTrace();
				}
				
				preds.put(child.getName(), inst);
			}
		}
		return preds;
		
	}

	private Map<String, CMUXML> loadGoldXMLs(File goldDir, int winSize, int stepSize) {
		Map<String, CMUXML> gold = new HashMap<String, CMUXML>();
		int count = 0, total = 0;
		if (goldDir.isDirectory()) {
			for (File child : goldDir.listFiles()) {
				if (child.getName().startsWith(".")) {
					continue;  // Ignore hidden files
				}
	
				CMUXML parser = new CMUXML(child, winSize, stepSize);
				parser.parse();
				if (parser.verifyData()) count++;
				total++;
				gold.put(child.getName(), parser);
			}
		}
		System.out.println(count + "/" + total + " videos verified.");
		return gold;
	}
	

	public static void main(String[] args) {
		System.out.println("Start\n");
		
		PerFramePRMetric metric = new PerFramePRMetric();
		metric.evaluate(new File("/Users/trananh/Desktop/psi/input/gold-xml"), 
				new File("/Users/trananh/Desktop/psi/output"), 30, 5, true);
		
		System.out.println("\nEnd");
	}
}
